rm(list = ls())
#-------------------------------------------------------------------------------------
setwd('~/Dropbox/Research/compliance_blocking/replication/Simulations')
library(logr)
lf<-log_open("02_simulation_analysis.log")
#LOAD IN NECESSARY RESULTS:
load('Simulation Results/sims_results_all.Rdata')

library(ggplot2)
library(tidyverse)
ggMelody<-theme_minimal() + theme(plot.title = element_text(hjust = 0.5, size = 18, face = "bold"),
                             axis.text=element_text(size=14), #size = 12
                             axis.title = element_text(size = 14),
                             legend.text=element_text(size=12),
                             legend.title =element_text(size=12),
                             legend.position='bottom',
                             strip.text.x = element_text(size = 14, face='bold'),
                             strip.text.y = element_text(size = 14, face='bold'),
                             plot.subtitle=element_text(size = 14, hjust=0.5),
                             panel.grid.minor = element_blank())
theme_set(ggMelody)

#COLOR PALETTE (very important):
pl1 = c("darkorange2", "#ffc14d", "#3dbbf5", "#00d0ca", "#88d575")

log_print(sessionInfo())

recode<-function(df_plot){
    df_plot$Estimator[which(df_plot$Estimator=="iv")]<-"IV"
    df_plot$Estimator[which(df_plot$Estimator=="psw")]<-"PSW"
    df_plot$Estimator[which(df_plot$Estimator=="iv_block")]<-"IV (Block)"
    df_plot$Estimator[which(df_plot$Estimator=="bdim")]<-"Block DiM"
    df_plot$Estimator[which(df_plot$Estimator=="psw_block")]<-"PSW (Block)"
    df_plot$Estimator[which(df_plot$Estimator=="placebo")]<-"Placebo (CR)"
    df_plot$Estimator[which(df_plot$Estimator=="placebo_block")]<-"Placebo (Block)"
    #Drop psw_full:
    df_plot$Estimator<-factor(df_plot$Estimator, levels=c("IV", "IV (Block)",
        "Placebo (CR)", "Placebo (Block)", "Block DiM", "PSW", "PSW (Block)"))

    df_plot$block_var[df_plot$block_var == "X2"]<-"Scenario 1. Compliance"
    df_plot$block_var[df_plot$block_var == "X1, X2"]<-"Scenario 2. Compliance & Outcome"
    df_plot$block_var[df_plot$block_var == "X1"]<-"Scenario 3. Outcome"
    df_plot$block_var[df_plot$block_var == "X4"]<-"Scenario 4. Noise"
    return(df_plot)
}
#-------------------------------------------------------------------------------------------------------
results_all = rbind(results_cov, results_comp, results_both)
df_se = results_all %>%
    group_by(block_var, n_sample) %>%
        summarize(
            iv=sd(iv),
            iv_block=sd(iv_block),
            placebo = sd(placebo_srs),
            placebo_block = sd(placebo_block),
            psw = sd(psw),
            psw_block = sd(psw_block),
            bdim = sd(block_dim)
        )
#----------------------------------------------------------------------------------------
#TABLES:
results_all$blocking = NA
results_all$blocking[results_all$block_var == "X1"]="Scenario 3: Block on Outcome-Related Variables"
results_all$blocking[results_all$block_var == "X2"]="Scenario 1: Block on Compliance-Related Variables"
results_all$blocking[results_all$block_var == "X1, X2"]="Scenario 2: Block on Compliance and Outcome-Related Variables"
df_mse <- results_all %>%
            group_by(blocking, n_sample) %>%
                summarize(
                    at = mean((AT - CACE)^2),
                    iv=mean((iv- CACE)^2),
                    iv_block=mean((iv_block - CACE)^2),
                    bdim = mean((block_dim - CACE)^2),
                    psw = mean((psw - CACE)^2),
                    psw_block = mean((psw_block - CACE)^2),
                    placebo = mean((placebo_srs - CACE)^2),
                    placebo_block = mean((placebo_block - CACE)^2)
                )


df_bias <- results_all %>%
            group_by(blocking, n_sample) %>%
                summarize(
                    at = mean(AT) - mean(CACE),
                    iv = mean(iv) - mean(CACE),
                    iv_block = mean(iv_block) - mean(CACE),
                    bdim = mean(block_dim) - mean(CACE),
                    psw = mean(psw) - mean(CACE),
                    psw_block = mean(psw_block) - mean(CACE),
                    placebo = mean(placebo_srs) - mean(CACE),
                    placebo_block = mean(placebo_block) -  mean(CACE),

                )
#TABLE 1
log_print(xtable::xtable(data.frame(df_mse[,-c(1, 9:10)], df_bias[,-c(1:2, 9:10)])), include.rownames=FALSE)

#------------------------------------------------------------------------------------------
#PI VIOLATION:
load('Simulation Results/sims_pi_violation.Rdata')
df_mse_pi <- results_pi %>%
            group_by(block_var, beta) %>%
                summarize(
                    #at = mean((AT - CACE)^2),
                    iv=mean((iv-CACE)^2),
                    iv_block=mean((iv_block - CACE)^2),
                    placebo = mean((placebo_srs - CACE)^2),
                    placebo_block = mean((placebo_block - CACE)^2),
                    psw = mean((psw-CACE)^2),
                    psw_block = mean((psw_block - CACE)^2),
                    bdim = mean((block_dim - CACE)^2)
                )


df_var_pi <- results_pi %>%
            group_by(block_var, beta) %>%
                summarize(
                    #at = mean((AT - CACE)^2),
                    iv=var(iv),
                    iv_block=var(iv_block),
                    placebo = var(placebo_srs),
                    placebo_block = var(placebo_block),
                    psw = var(psw),
                    psw_block = var(psw_block),
                    bdim = var(block_dim)
                )

#BIAS:
results_pi %>%
            group_by(block_var, beta) %>%
                summarize(
                    #at = mean((AT - CACE)^2),
                    iv=mean(iv) - mean(CACE),
                    iv_block=mean(iv_block) - mean(CACE),
                    psw = mean(psw) - mean(CACE),
                    psw_block = mean(psw_block) - mean(CACE),
                    bdim = mean(block_dim) - mean(CACE)
                )

df_plot_pi = rbind(
    data.frame(df_mse_pi %>% gather(key = Estimator, value = Value, -c(block_var, beta)), Measure="Bias^2"),
    data.frame(df_var_pi %>% gather(key = Estimator, value = Value, -c(block_var, beta)), Measure="Variance")
    )

df_plot_pi = recode(df_plot_pi)

#FIGURE 2
df_plot_pi %>%
filter(!Estimator %in% c("Placebo (CR)", "Placebo (Block)")) %>%
ggplot(aes(x = as.factor(beta), y =Value, fill=Estimator, alpha=Measure, group=Estimator))+
geom_bar(stat = 'identity', position='dodge', width=0.8)+
facet_wrap(~block_var)+
scale_alpha_manual(values = c(0.4, 1))+
ggtitle("MSE under Violations of Principal Ignorability")+
xlab("Correlation with Target Blocking Variable") + ylab("Mean Squared Error")+
scale_fill_manual(values=pl1)
ggsave("../Figures/fig2_mse_breakdown_pi.pdf", width=12, height=4.5)
#------------------------------------------------------------------------------------------
#VARYING COMPLIANCE PROBABILITY:
load('Simulation Results/sims_comp_prob.Rdata')

df_se_prob = results_comp_prob %>%
    group_by(block_var, compliance_rate) %>%
        summarize(
            iv=sd(iv),
            iv_block=sd(iv_block),
            placebo = sd(placebo_srs),
            placebo_block = sd(placebo_block),
            psw = sd(psw),
            psw_block = sd(psw_block),
            bdim = sd(block_dim)
        )

df_plot_prob = df_se_prob %>%
gather(key = Estimator, value=SE, -c(block_var, compliance_rate)) %>% recode()

#FIGURE 3:
df_plot_prob %>% filter(!Estimator %in% c("Placebo (CR)", "Placebo (Block)")) %>%
    ggplot(aes(x = as.factor(compliance_rate), y = SE,
        fill=Estimator, group=Estimator)) + facet_wrap(~block_var) +
    geom_bar(stat='identity', position='dodge', width=0.75)+
    ggtitle("Standard Error Comparison Across Compliance Rates")+ xlab("Compliance Rate") +
    ylab("Empirical Standard Error") + scale_fill_manual(values = pl1)
ggsave('../Figures/fig3_standard_error_compl.pdf', width=10.5, height=4)


log_close()